# 左键前景，右键背景
import cv2
import numpy as np

#绘制前景/背景标识线标志
drawing = False

# 定义GrabCut类，作用是设置一些参数
class GrabCut:
    def __init__(self, t_img):
        self.img = t_img
        self.img_raw = img.copy()
        self.img_width = img.shape[0]
        self.img_height = img.shape[1]
        self.img_show = self.img.copy()
        self.img_gc = self.img.copy()
        self.img_gc = cv2.GaussianBlur(self.img_gc, (3, 3), 0)
        self.lb_up = False
        self.rb_up = False
        self.lb_down = False
        self.rb_down = False
        self.mask = np.full(self.img.shape[:2], 2, dtype=np.uint8)
        self.firt_choose = True

# 鼠标操作的的回调函数
def mouse_event(event, x, y, flags, param):
    global drawing, last_point, start_point
    # 左键按下，开始标识前景
    if event == cv2.EVENT_LBUTTONDOWN:
        drawing = True
        # 设置鼠标按下的起始点
        last_point = (x, y)
        start_point = last_point
        param.lb_down = True
    # 右键按下，开始标识背景
    elif event == cv2.EVENT_RBUTTONDOWN:
        # 用户请先标识前景，否则无法分割
        if param.firt_choose:
            print("Please select foreground first!")
            return
        drawing = True
        last_point = (x, y)
        start_point = last_point
        param.rb_down = True
    # 鼠标移动，绘制标识前景和背景的线
    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing:
            # 鼠标左键按下的绘制
            if param.lb_down:
                cv2.line(param.img_show, last_point, (x,y), (0, 0, 255), 10, -1)
                cv2.rectangle(param.mask, last_point, (x, y), 1, -1, 4)
            # 鼠标右键按下的绘制
            if param.rb_down:
                cv2.line(param.img_show, last_point, (x, y), (255, 0, 0), 10, -1)
                cv2.rectangle(param.mask, last_point, (x, y), 0, -1, 4)
            last_point = (x, y)
    # 左键释放，结束标识前景
    elif event == cv2.EVENT_LBUTTONUP:
        drawing = False
        param.lb_up = True
        param.lb_down = False
        cv2.line(param.img_show, last_point, (x,y), (0, 0, 255), 2, -1)
        # 如果第一次标识，切换状态
        if param.firt_choose:
            param.firt_choose = False
        cv2.rectangle(param.mask, last_point, (x,y), 1, -1, 4)
    # 右键释放，结束标识背景
    elif event == cv2.EVENT_RBUTTONUP:
        # 如果首先标识背景则不做处理
        if param.firt_choose:
            return
        drawing = False
        param.rb_up = True
        param.rb_down = False
        cv2.line(param.img_show, last_point, (x,y), (255, 0, 0), 2, -1)
        cv2.rectangle(param.mask, last_point, (x,y), 0, -1, 4)

#执行操作
def process(img):
    if img is None:
        print('Can not read image correct!')
        return
    g_img = GrabCut(img)

    cv2.namedWindow('image')
    # 定义鼠标的回调函数
    cv2.setMouseCallback('image', mouse_event, g_img)
    while (True):
        cv2.imshow('image', g_img.img_show)
        # 鼠标左键或者右键抬起时，按照标识执行Grabcut算法
        if g_img.lb_up or g_img.rb_up:
            g_img.lb_up = False
            g_img.rb_up = False
            # 背景model
            bgdModel = np.zeros((1, 65), np.float64)
            # 前景model
            fgdModel = np.zeros((1, 65), np.float64)

            rect = (1, 1, g_img.img.shape[1], g_img.img.shape[0])
            mask = g_img.mask
            g_img.img_gc = g_img.img.copy()
            #执行Grabcut算法
            cv2.grabCut(g_img.img_gc, mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_MASK)
            # 0和2做背景
            mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype('uint8')
            # 使用蒙板来获取前景区域
            g_img.img_gc = g_img.img_gc * mask2[:, :, np.newaxis]
            cv2.imshow('Grabcut_result', g_img.img_gc)
            key = cv2.waitKey(0)
            if key == 115:
                # 假设 g_img.img_gc 是前景提取后的图像（BGR）
                # mask2 是 GrabCut 结果，前景=1 或 255，背景=0
                # 为了安全统一成 0/255
                alpha = np.where(mask2 == 0, 0, 255).astype(np.uint8)
                # 构造 BGRA 图像
                b, g, r = cv2.split(g_img.img_gc)
                bgra = cv2.merge([b, g, r, alpha])
                # 保存为 PNG（支持透明）
                cv2.imwrite("target.png", bgra)

        # 按下ESC键退出
        if cv2.waitKey(20) == 27:
            break

if __name__ == '__main__':
    img = cv2.imread(r"C:\Users\Administrator\Desktop\OpenCV4_Book_code\qq.png")
    process(img)